import pickle
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go


def read_pickle_file(filename):
  with open(filename, 'rb') as file:
    return pickle.load(file)


def inject_dict_into_dataframe(meta, data):
  meta_len = data.__len__()
  for idx, (key, value) in enumerate(meta.items()):
    tmp_row = value
    layers, kernel_l, kernel_av, lr = [int(str) for str in key.split('_')]
    tmp_row['num_layers'] = layers
    tmp_row['kernel_size_l'] = kernel_l
    tmp_row['reception_field_span_av'] = kernel_av
    tmp_row['lr'] = lr
    data.loc[meta_len + idx] = tmp_row
  return data


data = None
for i in range(31, 35):
  meta = read_pickle_file(f'result{i}.pkl')

  if data is None:
    columns = []
    columns = ['num_layers', 'kernel_size_l', 'reception_field_span_av', 'lr']
    columns.extend(meta[list(meta.keys())[0]].keys())
    data = pd.DataFrame(columns=columns)

  data = inject_dict_into_dataframe(meta, data)

data['k_l_tot'] = (data['kernel_size_l'] - 1) * data['num_layers'] + 1
data['k_av_tot'] = (data['reception_field_span_av'] -
                    1) * data['num_layers'] + 1

fig = px.box(data, x='num_layers', y="f1_Overall", points="all")
fig.show()

fig = px.box(data, x='kernel_size_l', y="f1_Overall", points="all")
fig.show()

fig = px.box(data, x='reception_field_span_av', y="f1_Overall", points="all")
fig.show()

fig = go.Figure()
for layers in range(2, 8):
  fig.add_trace(
      go.Box(
          x=data[data['num_layers'] == layers]['kernel_size_l'],
          y=data[data['num_layers'] == layers]['f1_Overall'],
          name=f'num_layers_{layers}'))
fig.update_layout(
    xaxis_title='kernel_size_l',
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for layers in range(2, 8):
  fig.add_trace(
      go.Box(
          x=data[data['num_layers'] == layers]['reception_field_span_av'],
          y=data[data['num_layers'] == layers]['f1_Overall'],
          name=f'num_layers_{layers}'))
fig.update_layout(
    xaxis_title='reception_field_span_av',
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for kernel_timing in range(17, 32 + 1, 5):
  fig.add_trace(
      go.Box(
          x=data[data['reception_field_span_av'] ==
                 kernel_timing]['num_layers'],
          y=data[data['reception_field_span_av'] == kernel_timing]
          ['f1_Overall'],
          name=f'reception_field_span_av_{kernel_timing}'))
fig.update_layout(
    xaxis_title='num_layers',
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for kernel_size in range(3, 20, 2):
  fig.add_trace(
      go.Box(
          x=data[data['kernel_size_l'] == kernel_size]['num_layers'],
          y=data[data['kernel_size_l'] == kernel_size]['f1_Overall'],
          name=f'kernel_size_l_{kernel_size}'))
fig.update_layout(
    xaxis_title='num_layers',
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for kernel_timing in range(17, 32 + 1, 5):
  fig.add_trace(
      go.Box(
          x=data[data['reception_field_span_av'] ==
                 kernel_timing]['kernel_size_l'],
          y=data[data['reception_field_span_av'] == kernel_timing]
          ['f1_Overall'],
          name=f'av_{kernel_timing}'))
fig.update_layout(
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

fig = go.Figure()
for kernel_size in range(3, 20, 2):
  fig.add_trace(
      go.Box(
          x=data[data['kernel_size_l'] == kernel_size]
          ['reception_field_span_av'],
          y=data[data['kernel_size_l'] == kernel_size]['f1_Overall'],
          name=f'l_{kernel_size}'))
fig.update_layout(
    yaxis_title='f1_Overall',
    boxmode='group'  # group together boxes of the different traces for each value of x
)
fig.show()

for layers in range(2, 8):
  fig = go.Figure()
  for kernel_size in range(3, 20, 2):
    # for kernel_timing in range(17, 32 + 1, 5):
    data_tmp = data[(data['num_layers'] == layers)
                    & (data['kernel_size_l'] == kernel_size)]
    fig.add_trace(
        go.Box(
            x=data_tmp['reception_field_span_av'],
            y=data_tmp['f1_Overall'],
            name=f'l_{layers}_kl_{kernel_size}'))
  fig.update_layout(
      yaxis_title='f1_Overall',
      boxmode='group'  # group together boxes of the different traces for each value of x
  )
  fig.show()

for layers in range(2, 8):
  fig = go.Figure()
  for kernel_timing in range(17, 32 + 1, 5):
    data_tmp = data[(data['num_layers'] == layers)
                    & (data['reception_field_span_av'] == kernel_timing)]
    fig.add_trace(
        go.Box(
            x=data_tmp['kernel_size_l'],
            y=data_tmp['f1_Overall'],
            name=f'l_{layers}_kav_{kernel_timing}'))
  fig.update_layout(
      yaxis_title='f1_Overall',
      boxmode='group'  # group together boxes of the different traces for each value of x
  )
  fig.show()

fig = px.density_heatmap(
    data,
    x='reception_field_span_av',
    y='num_layers',
    z="f1_Overall",
    nbinsx=4,
    nbinsy=6,
    color_continuous_scale=px.colors.sequential.Greys,
    histfunc='avg')
fig.show()

fig = px.density_heatmap(
    data,
    x='kernel_size_l',
    y='num_layers',
    z="f1_Overall",
    nbinsx=9,
    nbinsy=6,
    color_continuous_scale=px.colors.sequential.Greys,
    histfunc='avg')
fig.show()

fig = px.density_heatmap(
    data,
    x='kernel_size_l',
    y='reception_field_span_av',
    z="f1_Overall",
    nbinsx=9,
    nbinsy=6,
    color_continuous_scale=px.colors.sequential.Greys,
    histfunc='avg')
fig.show()

fig = px.density_heatmap(
    data,
    x="k_l_tot",
    y="k_av_tot",
    z="f1_Overall",
    nbinsx=9,
    nbinsy=6,
    color_continuous_scale=px.colors.sequential.Greys,
    histfunc='avg')
fig.show()

layer_list = list(range(2, 8))
for layers in layer_list:
  fig = px.density_heatmap(
      data[data['num_layers'] == layers],
      x='kernel_size_l',
      y='reception_field_span_av',
      z="f1_Overall",
      nbinsx=9,
      nbinsy=6,
      color_continuous_scale=px.colors.sequential.Greys,
      histfunc='avg')
  fig.show()

print(meta)